import bisect  # For efficient searching in sorted list
import csv
import os
from datetime import datetime, timezone
from typing import List, Tuple, Union

# --- Configuration ---
BASE_DATA_PATH = r"F:\personal\binance_klines"  # Base directory for all coin data

# Expected header columns - used to detect and skip header
EXPECTED_HEADER_FIELDS = [
    "open_time", "open", "high", "low", "close", "volume",
    "close_time", "quote_volume", "count",
    "taker_buy_volume", "taker_buy_quote_volume", "ignore"
]


class KLine:
    """
    Represents a single candlestick (Kline) data point.
    Timestamps are stored as milliseconds since epoch (UTC).
    """

    def __init__(self, data_row: List[str]):
        try:
            self.open_time: int = int(data_row[0])
            self.open: float = float(data_row[1])
            self.high: float = float(data_row[2])
            self.low: float = float(data_row[3])
            self.close: float = float(data_row[4])
            self.volume: float = float(data_row[5])
            self.close_time: int = int(data_row[6])
            self.quote_volume: float = float(data_row[7])
            self.count: int = int(data_row[8])  # Number of trades
            self.taker_buy_volume: float = float(data_row[9])
            self.taker_buy_quote_volume: float = float(data_row[10])
            self.ignore: str = data_row[11]  # Or float(data_row[11]) if always numeric
        except (IndexError, ValueError) as e:
            raise ValueError(f"Error parsing KLine data row: {data_row}. Error: {e}")

    @property
    def open_datetime_utc(self) -> datetime:
        return datetime.fromtimestamp(self.open_time / 1000, tz=timezone.utc)

    @property
    def close_datetime_utc(self) -> datetime:
        return datetime.fromtimestamp(self.close_time / 1000, tz=timezone.utc)

    def __repr__(self) -> str:
        return (f"KLine(OpenTime={self.open_datetime_utc.strftime('%Y-%m-%d %H:%M:%S')}, "
                f"O={self.open}, H={self.high}, L={self.low}, C={self.close}, V={self.volume})")

    # Allow comparison based on open_time for sorting
    def __lt__(self, other) -> bool:
        if isinstance(other, KLine):
            return self.open_time < other.open_time
        return NotImplemented


class KLines:
    """
    Manages a collection of KLine objects, loaded from CSV files for a specific
    coin, date_type, and date range.
    """

    def __init__(self, coin: str, date_type: str, start_date_str: str, end_date_str: str, base_data_path: str):
        """
        Initializes KLines and loads data.
        :param coin: Trading pair symbol (e.g., "ETHUSDT").
        :param date_type: Time interval (e.g., "1m", "1h", "1d").
        :param start_date_str: Start month in "YYYY-MM" format (inclusive).
        :param end_date_str: End month in "YYYY-MM" format (inclusive).
        """
        self.coin = coin
        self.date_type = date_type
        self.start_date_str = start_date_str
        self.end_date_str = end_date_str
        self.all_klines: List[KLine] = []  # Stores all loaded KLine objects, sorted by open_time
        self.data_path_prefix = os.path.join(base_data_path, self.coin, self.date_type)

        self._load_data()

    def _parse_year_month(self, date_str: str) -> Tuple[int, int]:
        try:
            year, month = map(int, date_str.split('-'))
            if not (1 <= month <= 12):
                raise ValueError("Month out of range")
            return year, month
        except ValueError:
            raise ValueError(f"Invalid date string format: '{date_str}'. Expected YYYY-MM.")

    def _load_data(self):
        """Loads KLine data from CSV files within the specified date range."""
        print(f"Loading data for {self.coin} ({self.date_type}) from {self.start_date_str} to {self.end_date_str}")
        print(f"Data path prefix: {self.data_path_prefix}")

        if not os.path.isdir(self.data_path_prefix):
            print(f"Warning: Directory not found: {self.data_path_prefix}. No data will be loaded.")
            return

        start_year, start_month = self._parse_year_month(self.start_date_str)
        end_year, end_month = self._parse_year_month(self.end_date_str)

        current_year = start_year
        current_month = start_month

        loaded_klines_temp: List[KLine] = []

        while current_year < end_year or \
                (current_year == end_year and current_month <= end_month):

            month_str = f"{current_month:02d}"
            csv_filename = f"{self.coin}-{self.date_type}-{current_year}-{month_str}.csv"
            csv_filepath = os.path.join(self.data_path_prefix, csv_filename)

            if os.path.exists(csv_filepath):
                print(f"  Processing file: {csv_filepath}")
                try:
                    with open(csv_filepath, 'r', newline='') as f:
                        reader = csv.reader(f)
                        is_first_row = True
                        for row_num, row in enumerate(reader):
                            if not row:  # Skip empty rows
                                continue
                            if is_first_row:
                                is_first_row = False
                                # Check if the first row is a header
                                # A simple check: if the first element of the row (lowercased)
                                # matches the first element of our expected header.
                                if row[0].strip().lower() == EXPECTED_HEADER_FIELDS[0].lower():
                                    print(f"    Header detected and skipped in {csv_filename}")
                                    continue
                                # else, it's data, process it

                            try:
                                kline_obj = KLine(row)
                                loaded_klines_temp.append(kline_obj)
                            except ValueError as e:
                                print(f"    Warning: Skipping malformed row {row_num + 1} in {csv_filename}: {e}")

                except Exception as e:
                    print(f"  Error reading file {csv_filepath}: {e}")
            else:
                print(f"  Warning: File not found: {csv_filepath}")

            # Increment month/year
            current_month += 1
            if current_month > 12:
                current_month = 1
                current_year += 1

        # Sort all loaded KLines by open_time
        self.all_klines = sorted(loaded_klines_temp)
        print(f"Finished loading. Total KLines loaded: {len(self.all_klines)}")

    @staticmethod
    def _to_utc_timestamp_ms(dt_input: Union[str, datetime]) -> int:
        """Converts ISO datetime string or datetime object to UTC milliseconds timestamp."""
        if isinstance(dt_input, str):
            dt_obj = datetime.fromisoformat(dt_input.replace("Z", "+00:00"))
        elif isinstance(dt_input, datetime):
            dt_obj = dt_input
        else:
            raise TypeError("Input must be an ISO datetime string or a datetime object.")

        # If naive, assume UTC. If aware, convert to UTC.
        if dt_obj.tzinfo is None or dt_obj.tzinfo.utcoffset(dt_obj) is None:
            dt_obj_utc = dt_obj.replace(tzinfo=timezone.utc)
        else:
            dt_obj_utc = dt_obj.astimezone(timezone.utc)

        return int(dt_obj_utc.timestamp() * 1000)

    def get_klines_by_range(self, start_datetime_iso: str, end_datetime_iso: str) -> List[KLine]:
        """
        Gets KLine objects within a specified datetime range.
        :param start_datetime_iso: Start datetime in ISO format (e.g., "2023-01-15T10:00:00Z" or "2023-01-15 10:00:00").
                                   Assumed UTC if no timezone info.
        :param end_datetime_iso: End datetime in ISO format (inclusive for open_time).
        :return: A list of KLine objects, sorted by open_time.
        """
        if not self.all_klines:
            return []

        try:
            start_ts_ms = self._to_utc_timestamp_ms(start_datetime_iso)
            end_ts_ms = self._to_utc_timestamp_ms(end_datetime_iso)
        except ValueError as e:
            print(f"Error parsing datetime strings: {e}")
            return []

        if start_ts_ms > end_ts_ms:
            print("Warning: Start datetime is after end datetime in get_klines_by_range.")
            return []

        # Since self.all_klines is sorted, we can use bisect for efficiency
        # Find the insertion point for start_ts_ms
        # We want klines where kline.open_time >= start_ts_ms
        dummy_start_kline = KLine([str(start_ts_ms)] + ['0'] * 11)  # Dummy for bisect_left
        start_index = bisect.bisect_left(self.all_klines, dummy_start_kline)

        result: List[KLine] = []
        for i in range(start_index, len(self.all_klines)):
            kline = self.all_klines[i]
            if kline.open_time <= end_ts_ms:  # Klines whose open_time is within the range
                result.append(kline)
            else:  # Past the end_datetime_iso
                break
        return result

    def get_klines_around_datetime(self, target_datetime_iso: str, count: int, direction: str = "forward") -> List[KLine]:
        """
        Gets a specified number of KLine objects immediately before or after a target datetime.
        :param target_datetime_iso: The reference datetime in ISO format.
        :param count: The number of KLine objects to retrieve.
        :param direction: "forward" to get KLines after (or at) the target,
                          "backward" to get KLines before the target.
        :return: A list of KLine objects, sorted by open_time (chronologically).
        """
        if not self.all_klines or count <= 0:
            return []

        try:
            target_ts_ms = self._to_utc_timestamp_ms(target_datetime_iso)
        except ValueError as e:
            print(f"Error parsing target datetime string: {e}")
            return []

        # Find the index of the KLine at or immediately after target_ts_ms
        # Create a dummy KLine for bisect.bisect_left
        # We use open_time for comparison. The other fields are arbitrary but must be valid types.
        dummy_target_kline = KLine([str(target_ts_ms)] + ['0'] * 11)

        # bisect_left finds insertion point to maintain order.
        # If a KLine has open_time == target_ts_ms, idx will point to it.
        # If no exact match, idx points to the first KLine with open_time > target_ts_ms.
        idx = bisect.bisect_left(self.all_klines, dummy_target_kline)

        result: List[KLine] = []
        if direction == "forward":
            # We want 'count' klines starting from index 'idx'
            # The kline at self.all_klines[idx] has open_time >= target_ts_ms
            end_slice = min(idx + count, len(self.all_klines))
            result = self.all_klines[idx:end_slice]
        elif direction == "backward":
            # We want 'count' klines ending *before* index 'idx' (or at idx-1 if target_ts_ms matches an existing
            # kline's open_time exactly, and we want data *strictly* before it) The kline at self.all_klines[idx-1]
            # has open_time < target_ts_ms (if idx > 0) Or, if target_ts_ms is exactly a kline's open_time,
            # self.all_klines[idx-1] is the one before it.

            # Correct starting point for backward slice is 'idx'.
            # We want elements up to, but not including, the one at 'idx'.
            # So, the slice should end at 'idx'.
            start_slice = max(0, idx - count)
            result = self.all_klines[start_slice:idx]  # This will be already sorted chronologically
        else:
            raise ValueError("Direction must be 'forward' or 'backward'.")

        return result

    def __len__(self) -> int:
        return len(self.all_klines)

    def __getitem__(self, index: int) -> KLine:
        return self.all_klines[index]


# --- Example Usage ---

if __name__ == "__main__":
    # Ensure the base directory structure exists for testing
    # F:\personal\binance_klines\ETHUSDT\1m\
    # And create some dummy CSV files for testing
    TEST_COIN = "ETHUSDT"
    TEST_DATATYPE = "1m"

    # --- Test with existing data (assuming you have 2023-01 and 2023-02 for ETHUSDT 1m) ---
    # If you don't have these specific dummy files, these tests might not behave as expected.
    # The KLines class will attempt to load what's available.
    # For robust testing of the class logic itself, isolated dummy files are better.
    # However, since you have real data, we can test against that.

    print("--- Initializing KLines object with some of your existing data range ---")
    try:
        # Let's assume you want to test with a small subset of your downloaded data
        # For example, Jan 2023 to Feb 2023 of ETHUSDT 1m data.
        # Ensure these files actually exist in your F:\personal\binance_klines\ETHUSDT\1m\ path
        EXISTING_START_MONTH = "2023-01"  # Use a month you know exists
        EXISTING_END_MONTH = "2023-02"  # Use another month you know exists

        klines_collection = KLines(
            coin=TEST_COIN,
            date_type=TEST_DATATYPE,
            start_date_str=EXISTING_START_MONTH,
            end_date_str=EXISTING_END_MONTH
        )

        print(f"\nTotal KLines loaded from your data: {len(klines_collection)}")
        if klines_collection.all_klines:
            print("First KLine from your data:", klines_collection.all_klines[0])
            print("Last KLine from your data:", klines_collection.all_klines[-1])
        else:
            print(
                f"Warning: No KLines loaded for {TEST_COIN} {TEST_DATATYPE} from {EXISTING_START_MONTH} to {EXISTING_END_MONTH}. Check paths and file existence.")

        # --- Testing get_klines_by_range (using your existing data) ---
        # Adjust these datetimes to fall within your EXISTING_START_MONTH and EXISTING_END_MONTH data
        if len(klines_collection) > 0:
            print("\n--- Testing get_klines_by_range (with your data) ---")
            # Example: Get klines from the first minute of your EXISTING_START_MONTH data
            # This requires knowing a valid timestamp range within your data.
            # For ETHUSDT-1m-2023-01.csv (if it starts at 2023-01-01 00:00:00 UTC)
            # open_time = 1672531200000
            range_start_dt = f"{EXISTING_START_MONTH.split('-')[0]}-{EXISTING_START_MONTH.split('-')[1]}-01T00:00:00Z"
            range_end_dt = f"{EXISTING_START_MONTH.split('-')[0]}-{EXISTING_START_MONTH.split('-')[1]}-01T00:01:00Z"  # Get first two 1m klines

            klines_in_range = klines_collection.get_klines_by_range(range_start_dt, range_end_dt)
            print(f"KLines from {range_start_dt} to {range_end_dt}:")
            for kline in klines_in_range:
                print(kline)
            # Add assertions based on your actual data if you know the expected output
            # e.g., assert len(klines_in_range) == 2 if you expect two klines in the first minute

            print("\n--- Testing get_klines_around_datetime (forward, with your data) ---")
            target_dt_fwd = range_start_dt  # Use the start of your data range
            klines_fwd = klines_collection.get_klines_around_datetime(target_dt_fwd, count=2, direction="forward")
            print(f"KLines forward from {target_dt_fwd} (count=2):")
            for kline in klines_fwd:
                print(kline)
            # Add assertions

            print("\n--- Testing get_klines_around_datetime (backward, with your data) ---")
            # Pick a time slightly after the start of your data to test backward
            # Example: 2023-01-01T00:02:00Z (assuming you have data for the first few minutes)
            target_dt_bwd = f"{EXISTING_START_MONTH.split('-')[0]}-{EXISTING_START_MONTH.split('-')[1]}-01T00:02:00Z"
            klines_bwd = klines_collection.get_klines_around_datetime(target_dt_bwd, count=2, direction="backward")
            print(f"KLines backward from {target_dt_bwd} (count=2):")
            for kline in klines_bwd:
                print(kline)
            # Add assertions

        # --- Test with non-existent month files (using a far future date) ---
        print("\n--- Test with non-existent month files (far future date) ---")
        NON_EXISTENT_MONTH_START = "2077-01"
        NON_EXISTENT_MONTH_END = "2077-01"

        klines_missing_month = KLines(
            coin=TEST_COIN,
            date_type=TEST_DATATYPE,
            start_date_str=NON_EXISTENT_MONTH_START,
            end_date_str=NON_EXISTENT_MONTH_END
        )
        print(f"Number of KLines loaded for non-existent month {NON_EXISTENT_MONTH_START}: {len(klines_missing_month)}")
        assert len(klines_missing_month) == 0, \
            f"Expected 0 klines for non-existent month {NON_EXISTENT_MONTH_START}, but got {len(klines_missing_month)}"

        print("\n--- Test with non-existent coin or date_type directory ---")
        # This test assumes that F:\personal\binance_klines\NONEXISTENTCOIN\ or F:\personal\binance_klines\ETHUSDT\nonexistent_type\ does NOT exist
        klines_missing_dir_coin = KLines(
            coin="NONEXISTENTCOIN",  # A coin you definitely don't have data for
            date_type=TEST_DATATYPE,
            start_date_str=EXISTING_START_MONTH,  # Use a valid date range
            end_date_str=EXISTING_END_MONTH
        )
        assert len(klines_missing_dir_coin) == 0, \
            f"Expected 0 klines for non-existent coin NONEXISTENTCOIN, but got {len(klines_missing_dir_coin)}"

        klines_missing_dir_type = KLines(
            coin=TEST_COIN,
            date_type="nonexistent_type",  # A date_type you definitely don't have
            start_date_str=EXISTING_START_MONTH,
            end_date_str=EXISTING_END_MONTH
        )
        assert len(klines_missing_dir_type) == 0, \
            f"Expected 0 klines for non-existent date_type nonexistent_type, but got {len(klines_missing_dir_type)}"


    except Exception as e:
        print(f"An error occurred during KLines operations: {e}")
        import traceback

        traceback.print_exc()
    finally:
        print("\nTest execution finished. Your original CSV files were not modified by these tests.")
